import os
import torch
import numpy as np
import torch.nn.functional as F
def get_cosine_similarity(data_a, data_b):
    tensor_a = torch.tensor(data_a, dtype=torch.float32).view(1, -1)
    tensor_b = torch.tensor(data_b, dtype=torch.float32).view(1, -1)
    tensor_a = torch.clamp(tensor_a, -100, 100)
    tensor_b = torch.clamp(tensor_b, -100, 100)
    assert tensor_a.shape[1] == tensor_b.shape[1]
    similarity_score = F.cosine_similarity(tensor_a, tensor_b).item()
    return similarity_score
def filter_files(dir, prefix='',postfix=''):
    files =[]
    for file in os.listdir(dir):
        if file.startswith(prefix) and file.endswith(postfix):
            files.append(os.path.join(dir,file))
    return files

def compare_npys(files1, files2):
    assert len(files1) == len(files2), "Files lists must have same length"
    files1 = sorted(files1)
    files2 = sorted(files2)
    for file1, file2 in zip(files1, files2):
        data1 = np.load(file1).clip(max=100.0,min=-100.0)
        data2 = np.load(file2).clip(max=100.0,min=-100.0)
        print(f'tensor:{os.path.basename(file1)}, similarity is: {get_cosine_similarity(data1, data2)}')




if __name__ == '__main__':

    q1 = np.fromfile('/home/adt/test/q1.bin')
    q2 = np.fromfile('/home/adt/test/q2.bin')
    sim = get_cosine_similarity(q1,q2)
    print(f'sim of q1 and q2 is: {sim}')

    key1 = np.fromfile('/home/adt/test/key1.bin')
    key2 = np.fromfile('/home/adt/test/key2.bin')
    sim = get_cosine_similarity(key1,key2)
    print(f'sim of key1 and key2 is: {sim}')

    v1 = np.fromfile('/home/adt/test/v1.bin')
    v2 = np.fromfile('/home/adt/test/v2.bin')
    sim = get_cosine_similarity(v1,v2)
    print(f'sim of v1 and v2 is: {sim}')

    kp1 = np.fromfile('/home/adt/test/kp1.bin')
    kp2 = np.fromfile('/home/adt/test/kp2.bin')
    sim = get_cosine_similarity(kp1,kp2)
    print(f'sim of kp1 and kp2 is: {sim}')

    qp1 = np.fromfile('/home/adt/test/qp1.bin')
    qp2 = np.fromfile('/home/adt/test/qp2.bin')
    sim = get_cosine_similarity(qp1,qp2)
    print(f'sim of qp1 and qp2 is: {sim}')

    m1 = np.fromfile('/home/adt/test/m1.bin')
    m2 = np.fromfile('/home/adt/test/m2.bin')
    sim = get_cosine_similarity(m1,m2)
    print(f'sim of m1 and m2 is: {sim}')


    mp1 = np.fromfile('/home/adt/test/mp1.bin')
    mp2 = np.fromfile('/home/adt/test/mp2.bin')
    sim = get_cosine_similarity(mp1,mp2)
    print(f'sim of mp1 and mp2 is: {sim}')

    l1 = np.fromfile('/home/adt/test/l1.bin')
    l2 = np.fromfile('/home/adt/test/l2.bin')
    sim = get_cosine_similarity(l1,l2)
    print(f'sim of l1 and l2 is: {sim}')

    hand1 = np.fromfile('/home/adt/test/hand1.bin')
    hand2 = np.fromfile('/home/adt/test/hand2.bin')
    sim = get_cosine_similarity(hand1,hand2)
    print(f'sim of hand1 and hand2 is: {sim}')

    o1 = np.fromfile('/home/adt/test/o1.bin')
    o2 = np.fromfile('/home/adt/test/o2.bin')
    sim = get_cosine_similarity(o1,o2)
    print(f'sim of o1 and o2 is: {sim}')

    rec_score1 = np.fromfile('/home/adt/test/rec_score1.bin')
    rec_score2 = np.fromfile('/home/adt/test/rec_score2.bin')
    sim = get_cosine_similarity(rec_score1,rec_score2)
    print(f'sim of rec_score1 and rec_score2 is: {sim}')

    memory_embedding1 = np.fromfile('/home/adt/test/memory_embedding1.bin')
    memory_embedding2 = np.fromfile('/home/adt/test/memory_embedding2.bin')
    sim = get_cosine_similarity(memory_embedding1,memory_embedding2)
    print(f'sim of memory_embedding1 and memory_embedding2 is: {sim}')


    PrepLidar_input_coord_idx0 = np.fromfile('/home/adt/test/tensors/cplus/PrepLidar_input_coord_idx0.npy')

    # files_cplus = filter_files('/home/adt/codes/python/width-former/WidthFormer/BEVDet/experiments/widthformer_slim_l_360/models/tensor/cplus',prefix='RPN_o',postfix='.npy')
    # files_py = filter_files('/home/adt/codes/python/width-former/WidthFormer/BEVDet/experiments/widthformer_slim_l_360/models/tensor/py',prefix='RPN_o',postfix='.npy')
    # compare_npys(files_cplus,files_py)

    pass